# -*- coding: utf-8 -*-

import akshare as ak
import pandas as pd
import os
import backtrader as bt
import datetime


class TestStrategy(bt.Strategy):
            
    def __init__(self):
        # 250日均线
        self.sma250 = bt.indicators.SimpleMovingAverage(self.data.close, period=250) 
        # 上浮20%, 下浮10%
        self.sma_up20 = bt.indicators.SimpleMovingAverage(self.data.close  * 1.2 , period=250)
        self.sma_down10 = bt.indicators.SimpleMovingAverage(self.data.close  * 0.9, period=250)
    
        # 设置绘图
        self.sma250.plotinfo.plotname = 'MA250'
        self.sma_up20.plotinfo.plotname = 'MA250 +20%'
        self.sma_down10.plotinfo.plotname = 'MA250 -10%'
        
        # 设置线条样式
        self.sma250.plotinfo.plot = True
        self.sma_up20.plotinfo.plot = True
        self.sma_down10.plotinfo.plot = True

    def next(self):
        pass

file_path = 'data/sh513050.csv'
if os.path.exists(file_path) and True:
    etf_df = pd.read_csv(file_path)
else:
    etf_df = ak.stock_zh_a_hist_tx(symbol='sh513050', start_date="20190101", end_date="20250223", adjust="qfq")
    etf_df.to_csv(file_path, encoding='utf-8')

# 将日期列转换为datetime格式并设置为索引
# etf_df['date'] = pd.to_datetime(etf_df['date'])
# etf_df.set_index('date', inplace=True)

etf_df = etf_df[['date', 'open', 'high', 'low', 'close','amount']]
etf_df.columns = ['date', 'open', 'high', 'low', 'close','amount']
etf_df.date = pd.to_datetime(etf_df['date'])
etf_df.index = etf_df.date
etf_df.sort_index(inplace=True)
etf_df.fillna(0.0,inplace=True)


# 读取和导入 dataframe 数据框 - 方式1
data = bt.feeds.PandasData(dataname=etf_df,
                        # 下面是数据表格默认包含的 7 个指标，取值对应指标在 daily_price.csv 中的列索引位置
                        datetime=0, # 告诉 GenericData， datetime 在 etf_df 文件的第0列
                        high=2,
                        low=3,
                        open=1,
                        close=4,
                        volume=5,
                        openinterest=-1) # 如果取值为 -1 , 告诉 GenericCSVData 该指标不存在

cerebro = bt.Cerebro()
st_date = datetime.datetime(2019,1,1)
ed_date = datetime.datetime(2025,2,23)
cerebro.adddata(data, name='sh513050')
cerebro.addstrategy(TestStrategy)
result = cerebro.run()
cerebro.plot()